"""8.8 通用mcp client集成各个组件开发-1"""
import asyncio
import json
from contextlib import AsyncExitStack
from typing import Any

from mcp import StdioServerParameters, stdio_client, ClientSession, ListToolsResult
from mcp.client.sse import sse_client
from mcp.types import CallToolResult, ReadResourceResult, GetPromptResult

from mcp_model import MCPFunction


class ServerInfo:
    def __init__(self,
                 name: str,  #服务端名称
                 cmd: str | None = None,
                 args: list[str] | None = None,
                 env: dict[str, Any] | None = None,
                 url: str | None = None,
                 ):
        self.name = name
        self.cmd = cmd
        self.args = args
        self.env = env
        self.url = url
        self.transport_type = None # stdio,sse
        if url is not None:
            self.transport_type = "sse"
        else:
            self.transport_type = "stdio"
        self.session:ClientSession = None


class ServerManager:
    def __init__(self):
        self.servers = {}
        self.async_exit_stack = AsyncExitStack()
        self.mcp_functions = {}

    async def initialize(self):
        await self.load_mcp_json()
        await self.create_server_session()
        await self.fetch_all_function_calling()

    async def load_mcp_json(self):
        with open("mcp.json","r") as f:
            mcp_json = json.load(f)
            mcp_servers = mcp_json["mcpServers"]
            #初始化ServerInfo对象
            for key,server in mcp_servers.items():
                server_info = ServerInfo(
                    name=key,
                    cmd=server.get("command"),
                    args=server.get("args"),
                    env=server.get("env"),
                    url=server.get("url")

                )
                #保存到字典中
                self.servers[server_info.name] = server_info

    async def create_server_session(self):
        for server_name,server in self.servers.items():
            #获取连接服务类型：sse,stdio
            transport_type = server.transport_type
            if transport_type == "stdio":
                # 一、创建服务连接的参数
                server_parameters = StdioServerParameters(
                    command=server.cmd,
                    args=server.args,
                    env=server.env
                )

                # 二、创建stdio_client
                client = stdio_client(server=server_parameters)

            elif transport_type == "sse":

                # 创建sse_client
                client = sse_client(url=server.url)
            else:
                raise ValueError("不支持的服务类型")
            transport = await self.async_exit_stack.enter_async_context(client)
            read_stream, write_stream = transport
            # 三、创建会话client
            client_session = ClientSession(read_stream, write_stream)
            session: ClientSession = await self.async_exit_stack.enter_async_context(client_session)
            server.session = session
            # 四、初始化会话
            await server.session.initialize()
    #获取所有的列表信息
    async def fetch_all_function_calling(self):
        for server_name,server in self.servers.items():
            session = server.session
            print("=== server_name:",server_name)
            # 列出所有可用的工具
            print(" =================== wrapper list_tools ========================")
            try:
                list_tool: ListToolsResult =await session.list_tools()
                tools = list_tool.tools
            except Exception as e:
                print("list_tools error:", e)
                tools = []

            for tool in tools:
                function = MCPFunction(
                    name=tool.name,
                    function_name=tool.name,
                    description=tool.description,
                    input_schema=tool.inputSchema,
                    server_name=server_name,
                    #类型
                    function_type="tool"
                )
                print(" tool--> ",function)
                self.mcp_functions[tool.name] = function
                #资源列表
                print(" =================== wrapper list_resources ========================")
                try:
                    list_resource = await session.list_resources()
                    resources = list_resource.resources
                except Exception as e:
                    print("list_resources error:", e)
                    resources = []


                for resource in resources:
                   # 去掉空格
                   resource_name = resource.name.replace(" ", "")
                   function = MCPFunction(
                       name=resource_name,
                       function_name=resource_name,
                       description=resource.description,
                       uri= resource.uri,
                       server_name=server_name,
                       #类型
                       function_type="resource"
                   )
                   print(" resource--> ", function)
                   self.mcp_functions[resource_name] = function
                #提示列表
                print(" =================== wrapper list_prompts ========================")
                try:
                    list_prompt = await session.list_prompts()
                    prompts = list_prompt.prompts
                except Exception as e:
                    print("list_prompts error:", e)
                    prompts = []

                for prompt in prompts:
                    function = MCPFunction(
                        name=prompt.name,
                        function_name=prompt.name,
                        description=prompt.description,
                        arguments= prompt.arguments,
                        server_name=server_name,
                        #类型
                        function_type="prompt"
                    )
                    print(" prompt--> ", function)
                    self.mcp_functions[prompt.name] = function
                    print(" =================== wrapper list_prompts end ========================")

    #执行Function calling
    async def execute_call(self,function_name,arguments:dict[str,Any]):
        mcp_function = self.mcp_functions.get(function_name)
        if mcp_function is None:
            raise ValueError(f"Function {function_name} not found")
        session:ClientSession = self.servers[mcp_function.server_name].session
        #根据方法的类型执行function calling
        function_type = mcp_function.function_type
        if function_type == "tool":
            # 调用工具
            print(" =================== wrapper call_tool ========================")
            try:
                result:CallToolResult = await session.call_tool(
                    name=mcp_function.name,
                    arguments=arguments
                )
                #获取文本结果
                text = result.content[0].text
                print("执行工具类tool结果:", result)
                return text
            except Exception as e:
                print("call_tool error:", e)
        #执行resource资源类方法
        elif function_type == "resource":
            # 调用资源
            print(" =================== wrapper call_resource ========================")
            try:
                resource_result:ReadResourceResult = await session.read_resource(mcp_function.uri)
                text = resource_result.contents[0].text
                print("执行资源类resource结果:", resource_result)
                return text


            except Exception as e:
                print("read_resource error:", e)
        # 执行prompt提示类方法
        elif function_type == "prompt":
            # 调用prompt
            print(" =================== wrapper call_prompt ========================")
            try:
                prompt_result:GetPromptResult = await session.get_prompt(
                    name=mcp_function.name,
                    arguments=arguments
                )
                text = prompt_result.messages[0].content.text
                print("执行提示类prompt结果:", prompt_result)
                return text

            except Exception as e:
                print("get_prompt error:", e)
    #关闭资源
    async def aclose(self):
        await self.async_exit_stack.aclose()
async def test():
    servicemanager = ServerManager()
    try:
        await servicemanager.initialize()
        result =await servicemanager.execute_call(
            function_name="list_tables",
            arguments={}
        )
        print("工具执行的结果:",result)
    except Exception as e:
        print("初始化失败:",e)
        return
    finally:
        await servicemanager.aclose()

if __name__ == "__main__":
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    loop.run_until_complete(test())
